import albumentations as A
from albumentations.pytorch import ToTensorV2
import random
import cv2
import os
import numpy as np
aug_transform = A.Compose([
    A.HorizontalFlip(p=0.5),
    A.VerticalFlip(p=0.5),
    A.RandomRotate90(p=0.7),
    A.Downscale(scale_min=0.6,scale_max=0.8,p=0.1),
    A.CoarseDropout(max_holes=3,max_width=3,max_height=3,min_holes=1,p=0.1),
    A.Transpose(),
    A.RandomResizedCrop(height=256,width=256,scale=(0.6,1),p=0.1),
    A.RandomBrightnessContrast(p=0.1),
])



def resize_trans(image,mask,size):
    resize_fun = A.Resize(width=size[0],height=size[1])
    trans_data = resize_fun(image=image,mask=mask)
    return trans_data['image'],trans_data['mask']


def random_concat_image(name_list,data_path,image_size = 256):
    ## 在所有的数据集中随机选取4张数据进行拼接。
    image_name = random.sample(name_list,4)
    ## 四张图片是不同的尺寸。随机生成交点位置的坐标
    coor = [random.randint(image_size//4,image_size//4*3),random.randint(image_size//4,image_size//4*3)]
    image_dir_list = [os.path.join(data_path,p) for p in image_name]
    image = [cv2.imread(p,-1) for p in image_dir_list]
    label = [cv2.imread(p.replace('tif','png'))[:,:,0] for p in image_dir_list]
    ## 根据交点的坐标来计算四张数据的尺寸。
    resize_list = [coor,[image_size-coor[0],coor[1]],[coor[0],image_size-coor[1]],[image_size-coor[0],image_size-coor[1]]]
    image0, label0 = resize_trans(image[0], label[0], resize_list[0])
    image1, label1 = resize_trans(image[1], label[1], resize_list[1])
    image2, label2 = resize_trans(image[2], label[2], resize_list[2])
    image3, label3 = resize_trans(image[3], label[3], resize_list[3])
    image = np.concatenate((np.concatenate((image0,image1),axis=1),np.concatenate((image2,image3),axis=1)),axis=0)
    label = np.concatenate((np.concatenate((label0,label1),axis=1),np.concatenate((label2,label3),axis=1)),axis=0)
    ## 注意label的取值范围需要进行转换
    label = label - 1
    return image,label


